import torch
import torch.utils.data as tud
import torch.nn as nn

from model_utils import AnomalyConfusionMatrix, INPUT_SIZE, create_dataloader, device, train_and_test_model
from mtad_gat_modules import (
    ConvLayer,
    FeatureAttentionLayer,
    TemporalAttentionLayer,
    GRULayer,
    Forecasting_Model,
    ReconstructionModel,
)


class MTAD_GAT(nn.Module):
    """ MTAD-GAT model class.

    :param n_features: Number of input features
    :param window_size: Length of the input sequence
    :param out_dim: Number of features to output
    :param kernel_size: size of kernel to use in the 1-D convolution
    :param feat_gat_embed_dim: embedding dimension (output dimension of linear transformation)
           in feat-oriented GAT layer
    :param time_gat_embed_dim: embedding dimension (output dimension of linear transformation)
           in time-oriented GAT layer
    :param use_gatv2: whether to use the modified attention mechanism of GATv2 instead of standard GAT
    :param gru_n_layers: number of layers in the GRU layer
    :param gru_hid_dim: hidden dimension in the GRU layer
    :param forecast_n_layers: number of layers in the FC-based Forecasting Model
    :param forecast_hid_dim: hidden dimension in the FC-based Forecasting Model
    :param recon_n_layers: number of layers in the GRU-based Reconstruction Model
    :param recon_hid_dim: hidden dimension in the GRU-based Reconstruction Model
    :param dropout: dropout rate
    :param alpha: negative slope used in the leaky rely activation function

    """

    def __init__(
            self,
            n_features,
            window_size,
            out_dim,
            kernel_size=7,
            feat_gat_embed_dim=None,
            time_gat_embed_dim=None,
            use_gatv2=True,
            gru_n_layers=1,
            gru_hid_dim=150,
            forecast_n_layers=1,
            forecast_hid_dim=150,
            recon_n_layers=1,
            recon_hid_dim=150,
            dropout=0.2,
            alpha=0.2
    ):
        super(MTAD_GAT, self).__init__()

        self.name = "MTAD_GAT"
        self.conv = ConvLayer(n_features, kernel_size)
        self.feature_gat = FeatureAttentionLayer(
            n_features, window_size, dropout, alpha, feat_gat_embed_dim, use_gatv2)
        self.temporal_gat = TemporalAttentionLayer(n_features, window_size, dropout, alpha, time_gat_embed_dim,
                                                   use_gatv2)
        self.gru = GRULayer(3 * n_features, gru_hid_dim, gru_n_layers, dropout)
        self.forecasting_model = Forecasting_Model(
            gru_hid_dim, forecast_hid_dim, out_dim, forecast_n_layers, dropout)
        self.recon_model = ReconstructionModel(window_size, gru_hid_dim, recon_hid_dim, out_dim, recon_n_layers,
                                               dropout)

    def forward(self, x):
        # x shape (b, n, k): b - batch size, n - window size, k - number of features

        x = self.conv(x)
        h_feat = self.feature_gat(x)
        h_temp = self.temporal_gat(x)

        h_cat = torch.cat([x, h_feat, h_temp], dim=2)  # (b, n, 3k)

        _, h_end = self.gru(h_cat)
        h_end = h_end.view(x.shape[0], -1)  # Hidden state for last timestamp

        predictions = self.forecasting_model(h_end)
        recons = self.recon_model(h_end)

        return predictions, recons


def compute_mtad_gat_loss(model: nn.Module, x: torch.Tensor, **kwargs):
    _, decoded = model(x)
    return nn.MSELoss().to(device)(x, decoded)


def get_optimizer(model):
    return torch.optim.Adam(model.parameters(), lr=1e-3)


def build_mtad_gat_model(input_size, **kwargs):
    model = MTAD_GAT(n_features=input_size,
                     window_size=INPUT_SIZE, out_dim=input_size).to(device)
    optimizer = get_optimizer(model)
    return model, optimizer


def train_and_test_mtad_gat_model(normal_dataloader: tud.DataLoader,
                                  attack_dataloader: tud.DataLoader,
                                  epochs: int,
                                  dataset_name: str,
                                  data_dir: str,
                                  continue_train: bool) -> AnomalyConfusionMatrix:
    return train_and_test_model(compute_loss=compute_mtad_gat_loss,
                                build_model=build_mtad_gat_model,
                                normal_dl=normal_dataloader,
                                attack_dl=attack_dataloader,
                                epochs=epochs,
                                model_name='MTAD_GAT',
                                continue_train=continue_train,
                                dataset=dataset_name,
                                data_dir=data_dir,
                                is_recur=True)


if __name__ == '__main__':
    normal_dl, attack_dl = create_dataloader(ds_path='./data/MSL/C-1',
                                             timemode=False,
                                             reverse_label=False)

    train_and_test_mtad_gat_model(normal_dataloader=normal_dl,
                                  attack_dataloader=attack_dl,
                                  epochs=1,
                                  dataset_name='MSL/C-1',
                                  continue_train=False)
